#include "../Species/Species3D.h"
#include "../PhysicalField/PhysicalField3D.h"
#include "../Field/Field3D.h"
#include "CellParticles.h"

Species3D::Species3D(PicParams* params_, ParallelMPI* pmpi_, int ispec) :
    Species(params_, pmpi_, ispec),
    //because PartBoundCond<dim> has no default constructor
    part_bound_cond(params_, pmpi_, ispec)
{
    init_particles();
}


Species3D::~Species3D()
{

}


void Species3D::init_particles()
{
    particles_.resize(params->dims_space[0]);
    for(int i = 0; i < params->dims_space[0]; i++)
    {
        particles_[i].resize(params->dims_space[1]);
        for(int j = 0; j < params->dims_space[1]; j++)
        {
            particles_[i][j].resize(params->dims_space[2]);
        }
    }

    vector<double> cell_min(3);

    for(int i = 0; i < params->dims_space[0]; i++)
    {
        for(int j = 0; j < params->dims_space[1]; j++)
        {
            for(int k = 0; k < params->dims_space[2]; k++)
            {
                particles_[i][j][k] = new CellParticles<3>();
                cell_min[0] = params->local_min[0] + i * params->cell_length[0];
                cell_min[1] = params->local_min[1] + j * params->cell_length[1];
                cell_min[2] = params->local_min[2] + k * params->cell_length[2];
                particles_[i][j][k]->initialize(params, species_param, cell_min, params->cell_length);
            }

        }
    }
}


void Species3D::move(PhysicalField* fields)
{
    //cout<<"move particle 1111 "<<endl;
    LocalFields E_cell_fields[2][2][2], B_cell_fields[2][2][2];
    Field3D<double>* Ex3D = static_cast<Field3D<double>*>(fields->Ex_);
    Field3D<double>* Ey3D = static_cast<Field3D<double>*>(fields->Ey_);
    Field3D<double>* Ez3D = static_cast<Field3D<double>*>(fields->Ez_);
    Field3D<double>* Bx3D = static_cast<Field3D<double>*>(fields->Bx_);
    Field3D<double>* By3D = static_cast<Field3D<double>*>(fields->By_);
    Field3D<double>* Bz3D = static_cast<Field3D<double>*>(fields->Bz_);

    B_cell_fields[0][0][0].x = params->externB[0];
    B_cell_fields[0][0][1].x = params->externB[0];
    B_cell_fields[0][1][0].x = params->externB[0];
    B_cell_fields[0][1][1].x = params->externB[0];
    B_cell_fields[1][0][0].x = params->externB[0];
    B_cell_fields[1][0][1].x = params->externB[0];
    B_cell_fields[1][1][0].x = params->externB[0];
    B_cell_fields[1][1][1].x = params->externB[0];

    B_cell_fields[0][0][0].y = params->externB[1];
    B_cell_fields[0][0][1].y = params->externB[1];
    B_cell_fields[0][1][0].y = params->externB[1];
    B_cell_fields[0][1][1].y = params->externB[1];
    B_cell_fields[1][0][0].y = params->externB[1];
    B_cell_fields[1][0][1].y = params->externB[1];
    B_cell_fields[1][1][0].y = params->externB[1];
    B_cell_fields[1][1][1].y = params->externB[1];

    B_cell_fields[0][0][0].z = params->externB[2];
    B_cell_fields[0][0][1].z = params->externB[2];
    B_cell_fields[0][1][0].z = params->externB[2];
    B_cell_fields[0][1][1].z = params->externB[2];
    B_cell_fields[1][0][0].z = params->externB[2];
    B_cell_fields[1][0][1].z = params->externB[2];
    B_cell_fields[1][1][0].z = params->externB[2];
    B_cell_fields[1][1][1].z = params->externB[2];


    for(int i = 0; i < params->dims_space[0]; i++)
    {
        for(int j = 0; j < params->dims_space[1]; j++)
        {
            for(int k = 0; k < params->dims_space[2]; k++)
            {
                int ii = i + params->oversize[0];
                int jj = j + params->oversize[1];
                int kk = k + params->oversize[2];

                E_cell_fields[0][0][0].x = (*Ex3D)(ii, jj, kk);
                E_cell_fields[0][0][1].x = (*Ex3D)(ii, jj, kk+1);
                E_cell_fields[0][1][0].x = (*Ex3D)(ii, jj+1, kk);
                E_cell_fields[0][1][1].x = (*Ex3D)(ii, jj+1, kk+1);
                E_cell_fields[1][0][0].x = (*Ex3D)(ii+1, jj, kk);
                E_cell_fields[1][0][1].x = (*Ex3D)(ii+1, jj, kk+1);
                E_cell_fields[1][1][0].x = (*Ex3D)(ii+1, jj+1, kk);
                E_cell_fields[1][1][1].x = (*Ex3D)(ii+1, jj+1, kk+1);

                E_cell_fields[0][0][0].y = (*Ey3D)(ii, jj, kk);
                E_cell_fields[0][0][1].y = (*Ey3D)(ii, jj, kk+1);
                E_cell_fields[0][1][0].y = (*Ey3D)(ii, jj+1, kk);
                E_cell_fields[0][1][1].y = (*Ey3D)(ii, jj+1, kk+1);
                E_cell_fields[1][0][0].y = (*Ey3D)(ii+1, jj, kk);
                E_cell_fields[1][0][1].y = (*Ey3D)(ii+1, jj, kk+1);
                E_cell_fields[1][1][0].y = (*Ey3D)(ii+1, jj+1, kk);
                E_cell_fields[1][1][1].y = (*Ey3D)(ii+1, jj+1, kk+1);

                E_cell_fields[0][0][0].z = (*Ez3D)(ii, jj, kk);
                E_cell_fields[0][0][1].z = (*Ez3D)(ii, jj, kk+1);
                E_cell_fields[0][1][0].z = (*Ez3D)(ii, jj+1, kk);
                E_cell_fields[0][1][1].z = (*Ez3D)(ii, jj+1, kk+1);
                E_cell_fields[1][0][0].z = (*Ez3D)(ii+1, jj, kk);
                E_cell_fields[1][0][1].z = (*Ez3D)(ii+1, jj, kk+1);
                E_cell_fields[1][1][0].z = (*Ez3D)(ii+1, jj+1, kk);
                E_cell_fields[1][1][1].z = (*Ez3D)(ii+1, jj+1, kk+1);

                push(particles_[i][j][k], E_cell_fields, B_cell_fields);
            }

        }
    }

    //cout<<"paritcle number "<<particles_[2][2]->size();

    //cout<<"move particle 2222 "<<endl;
    //sort_particles();
    //cout<<"move particle 3333 "<<endl;
}

void Species3D::project(PhysicalField* fields)
{
    //cout<<"project begen ========="<<endl;
    Field3D<double>* rho3D = static_cast<Field3D<double>*>(fields->rho_s[species_number]);

    //correct_particles();

    for(int i = 0; i < params->dims_space[0]; i++)
    {
        for(int j = 0; j < params->dims_space[1]; j++)
        {
            for(int k = 0; k < params->dims_space[2]; k++)
            {
                double rhoQ_cell_fields[2][2][2];

                rhoQ_cell_fields[0][0][0] = 0.0;
                rhoQ_cell_fields[0][0][1] = 0.0;
                rhoQ_cell_fields[0][1][0] = 0.0;
                rhoQ_cell_fields[0][1][1] = 0.0;
                rhoQ_cell_fields[1][0][0] = 0.0;
                rhoQ_cell_fields[1][0][1] = 0.0;
                rhoQ_cell_fields[1][1][0] = 0.0;
                rhoQ_cell_fields[1][1][1] = 0.0;

                //if(particles_[i][j][k]->size() == 0) cout<<"particle number = 0 "<<i<<"  "<<j<<"  "<<k<<endl;
                for(auto &particle : particles_[i][j][k]->data_)
                {
                    double xpn = particle.position[0] * params->cell_length_inv[0];
                    double ypn = particle.position[1] * params->cell_length_inv[1];
                    double zpn = particle.position[2] * params->cell_length_inv[2];

                    int ip_ = floor(xpn);
                    int jp_ = floor(ypn);
                    int kp_ = floor(zpn);

                    //Declaration and calculation of the coefficient for interpolation
                    double delta;
                    double coeffxp[2], coeffyp[2], coeffzp[2];

                    delta   = xpn - (double)ip_;
                    coeffxp[0] = 1.0 - delta;
                    coeffxp[1] = delta;

                    delta   = ypn - (double)jp_;
                    coeffyp[0] = 1.0 - delta;
                    coeffyp[1] = delta;

                    delta   = zpn - (double)kp_;
                    coeffzp[0] = 1.0 - delta;
                    coeffzp[1] = delta;

                    for (int iloc=0 ; iloc<2 ; iloc++)
                    {
                        for (int jloc=0 ; jloc<2 ; jloc++)
                        {
                            for (int kloc=0 ; kloc<2 ; kloc++)
                            {
                                rhoQ_cell_fields[iloc][jloc][kloc] += coeffxp[iloc] * coeffyp[jloc] * coeffzp[kloc];
                            }

                        }
                    }
                }
                int ii = i + params->oversize[0];
                int jj = j + params->oversize[1];
                int kk = k + params->oversize[2];
                (*rho3D)(ii,jj,kk)      += rhoQ_cell_fields[0][0][0];
                (*rho3D)(ii,jj,kk+1)    += rhoQ_cell_fields[0][0][1];
                (*rho3D)(ii+1,jj,kk)    += rhoQ_cell_fields[1][0][0];
                (*rho3D)(ii+1,jj,kk+1)  += rhoQ_cell_fields[1][0][1];
                (*rho3D)(ii,jj+1,kk)    += rhoQ_cell_fields[0][1][0];
                (*rho3D)(ii,jj+1,kk+1)  += rhoQ_cell_fields[0][1][1];
                (*rho3D)(ii+1,jj+1,kk)  += rhoQ_cell_fields[1][1][0];
                (*rho3D)(ii+1,jj+1,kk+1)+= rhoQ_cell_fields[1][1][1];

            }

        }
    }

    //cout<<"project begen1111 ========="<<endl;
    for(int i = 0; i < params->dims_field[0]; i++)
    {
        for(int j = 0; j < params->dims_field[1]; j++)
        {
            for(int k = 0; k < params->dims_field[2]; k++)
            {
                (*rho3D)(i,j,k) *= species_param.weight;
                //if((*rho2D)(i,j) < 0.8e-18) cout<<i<<"  "<<j<<"  "<< (*rho2D)(i,j)<<endl;
            }

        }
    }

    //int imax = params->dims_field[0] - 1;
    //int jmax = params->dims_field[1] - 1;
    //cout<<(*rho2D)(0,0)<<"  "<<(*rho2D)(0,1)<<"  "<<(*rho2D)(1,1)<<"  "<<(*rho2D)(1,2)<<"  "<<(*rho2D)(2,2)<<"  "<<endl;
    //cout<<(*rho2D)(imax,jmax)<<"  "<<(*rho2D)(imax,jmax - 1)<<"  "<<(*rho2D)(imax - 1,jmax - 1)<<"  "<<(*rho2D)(imax - 1,jmax - 2)<<"  "<<(*rho2D)(imax - 2,jmax - 2)<<"  "<<endl;
    //cout<<"project begen22222 ========="<<endl;
}


void Species3D::push(CellParticles<3>* cell_particles, const LocalFields E_cell_fields[2][2][2], const LocalFields B_cell_fields[2][2][2])
{
    //double charge_over_mass = static_cast<double>(particles.charge(ipart))*one_over_mass_;
    double um0, um1, um2, up0, up1, up2, p0dot, p1dot, p2dot;
    double alpha, inv_det_t, t0, t1, t2, t0_square, t1_square, t2_square, s0, s1, s2;
    double t0t1, t1t2, t2t0;
    double p0sm, p1sm, p2sm;
    double dts2, charge_over_mass;
    LocalFields E_local, B_local;

    dts2 = params->dt * 0.5;
    charge_over_mass = species_param.charge / species_param.mass;
    alpha = charge_over_mass * dts2;

    //B is constant and uniform
    B_local.x = params->externB[0];
    B_local.y = params->externB[1];
    B_local.z = params->externB[2];

    for(auto &particle : cell_particles->data_)
    {
        interpolate(E_cell_fields, E_local, particle.position);
        //interpolate(B_cell_fields, B_local, particle.position);

        um0 = particle.velocity[0] + alpha * E_local.x;
        um1 = particle.velocity[1] + alpha * E_local.y;
        um2 = particle.velocity[2] + alpha * E_local.z;

        //Rotation in the magnetic field
        t0    = alpha * B_local.x;
        t1    = alpha * B_local.y;
        t2    = alpha * B_local.z;
        t0_square   = t0*t0;
        t1_square   = t1*t1;
        t2_square   = t2*t2;
        inv_det_t = 2.0/(1.0+t0_square+t1_square+t2_square);
        s0    = t0 * inv_det_t;
        s1    = t1 * inv_det_t;
        s2    = t2 * inv_det_t;

        p0dot = um0 + um1 * t2 - um2 * t1;
        p1dot = um1 + um2 * t0 - um0 * t2;
        p2dot = um2 + um0 * t1 - um1 * t0;

        up0 = um0 + p1dot * s2 - p2dot * s1;
        up1 = um1 + p2dot * s0 - p0dot * s2;
        up2 = um2 + p0dot * s1 - p1dot * s0;

        //half-acceleration in the electric field
        p0sm = up0 + charge_over_mass * E_local.x * dts2;
        p1sm = up1 + charge_over_mass * E_local.y * dts2;
        p2sm = up2 + charge_over_mass * E_local.z * dts2;

        particle.velocity[0] = p0sm;
        particle.velocity[1] = p1sm;
        particle.velocity[2] = p2sm;

        //move the particle
        for(int i = 0 ; i < 3 ; i++)
        {
            //particle.position_old[i]  = particle.position[i];
            particle.position[i] += params->dt * particle.velocity[i];
        }
    }

}



void Species3D::sort_particles()
{
    int new_index[3];

    psi_particles_.clear();
    exchange_particles_.clear();

    for(int i = 0; i < params->dims_space[0]; i++)
    {
        for(int j = 0; j < params->dims_space[1]; j++)
        {
            for(int k = 0; k < params->dims_space[2]; k++)
            {
                auto cell_particles = particles_[i][j][k];
                for(int i_particle = 0; i_particle < cell_particles->size(); i_particle++)
                {
                    //keep_particle = :  0 delete; 1 MPI change; 2 move to other cell; 3 still in the cell
                    auto &particle = cell_particles->data_[i_particle];
                    new_index[0] = floor( (particle.position[0] - params->local_min[0]) * params->cell_length_inv[0] );
                    new_index[1] = floor( (particle.position[1] - params->local_min[1]) * params->cell_length_inv[1] );
                    new_index[2] = floor( (particle.position[2] - params->local_min[2]) * params->cell_length_inv[2] );

                    int keep_particle = -1;
                    if(new_index[0] == i && new_index[1] == j && new_index[2] == k)
                    {
                        keep_particle = 3;
                    }
                    else if(new_index[0] >= 0 && new_index[0] < params->dims_space[0] && new_index[1] >= 0 && new_index[1] < params->dims_space[1] && new_index[2] >= 0 && new_index[2] < params->dims_space[2])
                    {
                        keep_particle = 2;
                    }
                    else
                    {
                        keep_particle = part_bound_cond.apply(particle);
                        //keep_particle = 0;
                    }

                    if (keep_particle == 0)
                    {
                        psi_particles_.push_back(particle);
                        cell_particles->add_particle_delete_list(i_particle);
                    }
                    else if(keep_particle == 1)
                    {
                        exchange_particles_.push_back(particle);
                        cell_particles->add_particle_delete_list(i_particle);
                    }
                    else if(keep_particle == 2)
                    {
                        particles_[new_index[0]][new_index[1]][new_index[2]]->add_new_particle(particle);
                        cell_particles->add_particle_delete_list(i_particle);
                    }
                    else if(keep_particle == -1)
                    {
                        log_error<<"error: keep_particle = 1";
                    }

                    //to be tested: how about performing psi_particles_.push_back(particle) and exchange_particles_.push_back(particle)
                    //              in another loop
                }
            }
        }
    }
}


void Species3D::correct_particles()
{
    for(int i = 0; i < params->dims_space[0]; i++)
    {
        for(int j = 0; j < params->dims_space[1]; j++)
        {
            for(int k = 0; k < params->dims_space[2]; k++)
            {
                auto cell_particles = particles_[i][j][k];
                cell_particles->delete_useless_particles();
            }

        }
    }
}

/*
void Species3D::interpolate(const LocalFields cell_fields[2][2][2], LocalFields& local_field, const double position[3])
{
    //Normalized particle position
    double xpn = position[0] * params->cell_length_inv[0];
    double ypn = position[1] * params->cell_length_inv[1];
    double zpn = position[2] * params->cell_length_inv[2];

    //Indexes of the central nodes
    int ip_ = floor(xpn);
    int jp_ = floor(ypn);
    int kp_ = floor(zpn);


    //Declaration and calculation of the coefficient for interpolation
    double delta;
    double coeffxp[2], coeffyp[2], coeffzp[2];

    delta   = xpn - (double)ip_;
    coeffxp[0] = 1.0 - delta;
    coeffxp[1] = delta;


    delta   = ypn - (double)jp_;
    coeffyp[0] = 1.0 - delta;
    coeffyp[1] = delta;

    delta   = zpn - (double)kp_;
    coeffzp[0] = 1.0 - delta;
    coeffzp[1] = delta;

    local_field.x = 0.0;
    local_field.y = 0.0;
    local_field.z = 0.0;

    for (int iloc=0 ; iloc<2 ; iloc++)
    {
        for (int jloc=0 ; jloc<2 ; jloc++)
        {
            for (int kloc=0 ; kloc<2 ; kloc++)
            {
                local_field.x += coeffxp[iloc] * coeffyp[jloc] * coeffzp[kloc] * cell_fields[iloc][jloc][kloc].x;
                local_field.y += coeffxp[iloc] * coeffyp[jloc] * coeffzp[kloc] * cell_fields[iloc][jloc][kloc].y;
                local_field.z += coeffxp[iloc] * coeffyp[jloc] * coeffzp[kloc] * cell_fields[iloc][jloc][kloc].z;
            }

        }
    }
}

*/
